package tree;

import org.junit.Test;

import java.util.Arrays;
import java.util.List;

public class CountCompleteTreeNodes222 {

    @Test
    public void test() {
        print(Arrays.asList(1,2,3,4,5,6));
        print(Arrays.asList());
        print(Arrays.asList(1));
        print(Arrays.asList(1,2,3,4,5,6,7));
    }

    void print(List<Integer> list) {
        TreeNode root = TreeNode.buildTree(list);
        int count = countNodes(root);
        System.out.format("tree: %s, count: %s\n", list, count);
    }

    // 先检查左右是不是满树, 是就公式计算, 不是就采用递归计算
    // 注意公式别搞错了
    public int countNodes(TreeNode root) {
        if (root == null) {
            return 0;
        }
        int leftDepth = getLeftDepth(root);
        int rightDepth = getRightDepth(root);
        if (leftDepth == rightDepth) {
            if (leftDepth == 0) {
                return 1;
            } else {
                return (int) Math.pow(2, leftDepth+1) - 1;
            }
        }
        return countNodes(root.left) + countNodes(root.right) + 1;
    }

    public int getLeftDepth(TreeNode node) {
        int depth = 0;
        while (node.left != null) {
            node = node.left;
            depth++;
        }
        return depth;
    }

    public int getRightDepth(TreeNode node) {
        int depth = 0;
        while (node.right != null) {
            node = node.right;
            depth++;
        }
        return depth;
    }
}
